#include "CryptoFile.h"
#include "Base64Encoder.h"
#include "SymmetricKeyMemory.h"

namespace kryptos { namespace cryptography {

CryptoFile::CryptoFile(const std::string& filename)
{
    // Try opening for reading first, if success then load it
    // if not, then its probably going to create a file so we
    // open it only with ios::out, it creats the file, the close it
    // and re-open it for both ios::in and ios::out
    m_stream.open(filename.c_str(), std::fstream::in | std::ios::binary);

    if (m_stream.good())
    {
        load();
        m_openForWritting = false;
    }
    else
    {
        m_stream.close();
        init(filename);
        m_openForWritting  = true;
        m_format           = (CryptoFile::CryptoFormat)-1;
    }
}

CryptoFile::CryptoFile(const std::string &filename, CryptoFormat format)
    : m_format(format)
{
    init(filename);
}

void CryptoFile::init(const std::string& filename)
{
    m_stream.open(filename.c_str(), std::fstream::out | std::ios::binary);

    if (m_stream.good())
    {
        m_stream.close();
        m_stream.open(filename.c_str(), std::fstream::out | std::fstream::in | std::ios::binary);
    }

    m_metadataBytes.clear();

    m_recipientsCount  = 0;
    m_generalHeader    = 0;
    m_cipherOffset     = 0;
    m_metadataOffset   = 0;
    m_recipientsSize   = 0;
    m_recipientsOffset = sizeof(m_format) + sizeof(m_recipientsOffset) + sizeof(m_recipientsCount) + sizeof(m_cipherOffset) + sizeof(m_metadataOffset);
    m_currentOffset    = m_recipientsOffset;       
}

CryptoFile::~CryptoFile()
{
    if (m_openForWritting)
    {
        /**
         *  Structure of file:
         *  ------------------
         *  - recipients offset
         *  - recipients count
         *  - metadata offset
         *  - cipherdata offset
         *  - data:
         *      - recipients
         *      - cipherdata
         *      - metadata
         **/

        // save metadata bytes at the end of the file
        if (m_format == K_OTPRSA01)
        {
            m_metadataOffset = m_stream.tellp();
            int metadatasize = m_cryptoMetadata.bytes().size();
            m_stream.write((const char*)&metadatasize, sizeof(int));
            m_stream.write((const char*)&m_cryptoMetadata.bytes()[0], metadatasize);
        }

        // write header data at the begining
        m_cipherOffset = m_recipientsOffset + m_recipientsSize;

        m_stream.seekp(m_generalHeader);
        m_stream.write((const char*)&m_format, sizeof(m_format));
        m_stream.write((const char*)&m_recipientsOffset, sizeof(m_recipientsOffset));
        m_stream.write((const char*)&m_recipientsCount, sizeof(m_recipientsCount));
        m_stream.write((const char*)&m_metadataOffset, sizeof(m_metadataOffset));
        m_stream.write((const char*)&m_cipherOffset, sizeof(m_cipherOffset));

        m_stream.flush();
    }

    m_stream.close();
}

void CryptoFile::addRecipient(const CryptoRecipient& recipient)
{
    m_openForWritting = true;
    std::vector<char> recipientBytes = recipient.bytes();
    m_stream.seekp(m_currentOffset);
    m_stream.write(&recipientBytes[0], recipientBytes.size());
    m_stream.flush();

    m_currentOffset = m_stream.tellp();
    m_recipientsSize += recipientBytes.size();
    m_recipientsCount++;
}

void CryptoFile::setMetadata(const CryptoMetadata& metadata)
{
    m_cryptoMetadata = metadata;
}

std::ostream& CryptoFile::outputStream()
{
    m_openForWritting = true;
    return m_stream;
}

std::istream& CryptoFile::inputStream()
{
    return m_stream;
}

unsigned long long CryptoFile::size()
{
    m_stream.seekg( 0, std::ios::end );
    std::streampos fsize = (unsigned long long)m_stream.tellg();
    m_stream.seekg(m_currentOffset);

    return fsize;
}

unsigned long long CryptoFile::cryptoSize()
{
    return (cryptoDataEnd() - m_cipherOffset);
}

unsigned long long CryptoFile::cryptoDataOffset()
{
    return m_cipherOffset;
}

unsigned long long CryptoFile::cryptoDataEnd()
{
    unsigned long long endCipherOffset = m_metadataOffset;
    if (endCipherOffset == 0)
    {
        m_stream.seekg(0, std::ios::end);
        endCipherOffset = (unsigned long long)m_stream.tellg();
        m_stream.seekg(m_currentOffset, std::ios::beg);
    }

    return endCipherOffset;
}

CryptoFile::CryptoFormat CryptoFile::format()
{
    return m_format;
}

std::vector<CryptoRecipient> CryptoFile::recipients()
{
    return m_recipients;
}

CryptoMetadata CryptoFile::metadata()
{
    return m_cryptoMetadata;
}

void CryptoFile::load()
{
    m_stream.read((char*)&m_format, sizeof(m_format));

    if (!m_stream.good() || (m_format != K_PK01 && m_format != K_OTPRSA01))
    {
        m_format = K_UNKNOWN;
    }
    else
    {
        m_stream.read((char*)&m_recipientsOffset, sizeof(m_recipientsOffset));
        m_stream.read((char*)&m_recipientsCount, sizeof(m_recipientsCount));
        m_stream.read((char*)&m_metadataOffset, sizeof(m_metadataOffset));
        m_stream.read((char*)&m_cipherOffset, sizeof(m_cipherOffset));

        for (int i=0; i<m_recipientsCount; ++i)
        {
            char id[40];
            std::vector<char> sessionKey(128);
            m_stream.read(id, 40);
            m_stream.read((char*)&sessionKey[0], 128);

            m_recipients.push_back(CryptoRecipient(std::string(id, 40), &sessionKey[0], sessionKey.size()));
        }

        m_currentOffset = m_cipherOffset;

        if (m_format == K_OTPRSA01)
        {
            m_stream.seekg(m_metadataOffset);

            int metadatasize;
            m_stream.read((char*)&metadatasize, sizeof(int));

            char *tmp = new char[metadatasize];
            m_stream.read(tmp, metadatasize);
            m_cryptoMetadata = CryptoMetadata((const char*)tmp, metadatasize);

            delete []tmp;
        }
    }
}

}}
